﻿// PPMatting_Staticlib.cpp : 定义静态库的函数。
//
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "PPMatting_StaticLib_cpu.h"

#ifdef WIN32
	const char sep = '\\';
#else
	const char sep = '/';
#endif

cv::Mat CpuInfer (const std::string &model_dir, const cv::Mat &image,
                  const std::string &background_file)
{
	auto model_file = model_dir + sep + "model.pdmodel";
	auto params_file = model_dir + sep + "model.pdiparams";
	auto config_file = model_dir + sep + "deploy.yaml";
	auto option = fastdeploy::RuntimeOption();
	option.UseCpu();
	auto model = fastdeploy::vision::matting::PPMatting (model_file, params_file,
	             config_file, option);
	cv::Mat vis_im;
	if (!model.Initialized()) {
		std::cerr << "Failed to initialize." << std::endl;
		return vis_im;
	}
	auto im = image;
	fastdeploy::vision::MattingResult res;
	if (!model.Predict (&im, &res)) {
		std::cerr << "Failed to predict." << std::endl;
		return vis_im;
	}
	
	if (!background_file.empty()) {
		auto bg = cv::imread (background_file);
		vis_im =
		    fastdeploy::vision::SwapBackground (im, bg, res);
	} else
		vis_im = fastdeploy::vision::VisMatting (im, res);
	return vis_im;
}

cv::Mat GpuInfer (const std::string &model_dir, const cv::Mat &image,
                  const std::string &background_file)
{
	auto model_file = model_dir + sep + "model.pdmodel";
	auto params_file = model_dir + sep + "model.pdiparams";
	auto config_file = model_dir + sep + "deploy.yaml";
	
	auto option = fastdeploy::RuntimeOption();
	option.UseGpu();
	
	option.UsePaddleInferBackend();
	auto model = fastdeploy::vision::matting::PPMatting (model_file, params_file,
	             config_file, option);
	cv::Mat vis_im;
	if (!model.Initialized()) {
		std::cerr << "Failed to initialize." << std::endl;
		return vis_im;
	}
	auto im = image;
	fastdeploy::vision::MattingResult res;
	
	if (!model.Predict (&im, &res)) {
		std::cerr << "Failed to predict." << std::endl;
		return vis_im;
	}
	
	
	if (!background_file.empty()) {
		auto bg = cv::imread (background_file);
		vis_im =
		    fastdeploy::vision::SwapBackground (im, bg, res);
	} else
		vis_im = fastdeploy::vision::VisMatting (im, res);
	return vis_im;
}

int infer_by_camera (const std::string &device, const std::string &model_dir,
                     const std::string &window_name,
                     const std::string &background_file)
{
	cv::VideoCapture cap;
	cap.open (0);
	if (!cap.isOpened()) {
		std::cout << "open camera failed!" << std::endl;
		return 0;
	}
	cv::namedWindow (window_name, 1);
	while (1) {
		time_t t_now = time (0);
		cv::Mat frame;
		cap >> frame;
		if (frame.empty())
			return 0;
		if (device == "gpu" or device == "GPU")
			cv::imshow (window_name, GpuInfer (model_dir, frame, background_file));
		else
			cv::imshow (window_name, CpuInfer (model_dir, frame, background_file));
		std::cout << "Matting此帧共消耗" << (time (0) - t_now) << "秒" << std::endl;
		if (cv::waitKey (30) >= 0)
			break;
	}
	cap.release();
	return 1;
}